/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package simple.utils.concurrent;

import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;

import java.io.Closeable;
import java.io.IOException;
import java.util.Collections;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author rehak
 */
public class ThreadWorker implements Closeable {


    private static final String DEF_NAME = "ThreadWorker";
    private final AtomicLong scheduledTasks = new AtomicLong(0);
    private final AtomicLong completedTasks = new AtomicLong(0);
    private final AlwaysBlockingQueue<Runnable> queue;
    private final ThreadPoolExecutor tpe;
    private Set<ListenableFuture<?>> running = Collections.newSetFromMap(new ConcurrentHashMap<ListenableFuture<?>, Boolean>());
    private final ListeningExecutorService executor;
    private final ExecutorService listenerExecutor = Executors.newSingleThreadExecutor();
    private final static AtomicLong no = new AtomicLong(0);
    private final AtomicLong threadNo = new AtomicLong(0);
    private final String name;

    public ThreadWorker(int maxThreads, int maxQueueSize) {
        this(maxThreads, maxQueueSize, DEF_NAME + no.getAndIncrement());
    }

    public ThreadWorker(int maxThreads, int maxQueueSize, String name) {
        this.name = name;
        this.queue = new AlwaysBlockingQueue<Runnable>(maxQueueSize);
        this.tpe = new ThreadPoolExecutor(maxThreads, maxThreads, 20, TimeUnit.SECONDS, queue);
        // customize thread factory
        final ThreadFactory tf = tpe.getThreadFactory();
        tpe.setThreadFactory(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread thread = tf.newThread(r);
                thread.setName(ThreadWorker.this.name + "-thread-" + threadNo.getAndIncrement());
                return thread;
            }

        });
        tpe.allowCoreThreadTimeOut(true);
        executor = MoreExecutors.listeningDecorator(tpe);
    }

    public int getQueueSize() {
        return queue.size();
    }

    public <E> ListenableFuture<E> submit(Callable<E> task) {
        scheduledTasks.incrementAndGet();
        final ListenableFuture<E> lf = executor.submit(task);
        running.add(lf);
        lf.addListener(new Runnable() {
            @Override
            public void run() {
                running.remove(lf);
                completedTasks.incrementAndGet();
            }
        }, listenerExecutor);
        return lf;
    }

    public long getCompletedTasks() {
        return completedTasks.get();
    }

    public long getScheduledTasks() {
        return scheduledTasks.get();
    }

    public int getRunningTasks() {
        return running.size();
    }

    public void interruptRunning() {
        for (ListenableFuture f : running) {
            f.cancel(true);
        }
    }

    public void blockUntilTasksAreDone() {
        while (running.size() > 0) {
            // do nothing
        }
    }

    public void waitForAll() throws InterruptedException {
        executor.shutdown();
        executor.awaitTermination(Integer.MAX_VALUE, TimeUnit.SECONDS);
        listenerExecutor.shutdown();
        listenerExecutor.awaitTermination(Integer.MAX_VALUE, TimeUnit.SECONDS);
    }

    @Override
    public void close() throws IOException {
        executor.shutdown();
        listenerExecutor.shutdown();
    }

}
